# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt

import sys
sys.path.append("./pyDmps") # add the pyDmps to import path
import cs
import dmp
import dmp_discrete

reload(dmp)
reload(dmp_discrete)

def single_trajectory_test(trajectory = None):
    plt.figure(2, figsize=(6, 4))
    num_bfs = [100, 10000]

    if trajectory is None:
        trajectory = np.sin(np.arange(0, 10, .01))

    for i, bfs in enumerate(num_bfs):
        # dmp = dmp_discrete.DMPs_discrete(dmps=1, bfs=bfs, dt=.01, goal=1, w=np.zeros((1, bfs)))
        dmp = dmp_discrete.DMPs_discrete(dmps=1, bfs=bfs, ay=np.array([25]), by=np.array([1]))

        dmp.imitate_path(trajectory)
        # dmp.goal = 1 # set the scale of movement

        scale_length = len(trajectory)
        tau_value = (1/dmp.dt)/scale_length
        y_track, dy_track, ddy_track = dmp.rollout(tau=tau_value)

        plt.figure(2)
        plt.plot(y_track)

    plt.plot(trajectory, 'r--', lw=2)
    plt.title('DMP imitate path')
    plt.xlabel('time (ms)')
    plt.ylabel('system trajectory')
    plt.legend(['%i BFs' % i for i in num_bfs], loc='lower right')
    plt.tight_layout()
    plt.show()

def two_trajectory_test():
    plt.figure(2, figsize=(6, 4))
    num_bfs = [10, 30, 50, 100, 10000]

    trajectory1 = np.sin(np.arange(0, 12, .01))
    trajectory2 = np.zeros(trajectory1.shape)
    trajectory2[int(len(trajectory2) / 2.):] = .5

    for i, bfs in enumerate(num_bfs):
        # dmp = dmp_discrete.DMPs_discrete(dmps=2, bfs=bfs, dt=.01, goal=1, w=np.zeros((1, bfs)))
        dmp = dmp_discrete.DMPs_discrete(dmps=2, bfs=bfs)

        dmp.imitate_path(y_des=np.array([trajectory1, trajectory2]))
        # change the scale of the movement
        # dmp.goal[0] = 1 # set the scale
        # dmp.goal[1] = 1 # set the scale

        scale_length = len(trajectory1)
        tau_value = (1/dmp.dt)/scale_length
        y_track, dy_track, ddy_track = dmp.rollout(tau=tau_value)

        plt.figure(2)
        plt.subplot(211)
        plt.plot(y_track[:, 0], lw=2)
        plt.subplot(212)
        plt.plot(y_track[:, 1], lw=2)

    plt.subplot(211)
    a = plt.plot(trajectory1, 'r--', lw=2)
    plt.title('DMP imitate trajectory')
    plt.xlabel('time (ms)')
    plt.ylabel('system trajectory')
    plt.legend([a[0]], ['desired trajectory'], loc='lower right')
    plt.subplot(212)
    b = plt.plot(trajectory2, 'r--', lw=2)
    plt.title('DMP imitate trajectory')
    plt.xlabel('time (ms)')
    plt.ylabel('system trajectory')
    plt.legend(['%i BFs' % i for i in num_bfs], loc='lower right')

    plt.tight_layout()
    plt.show()


def trajectory2D_test(trajectory = None):
    plt.figure(2, figsize=(6, 4))
    num_bfs = [100, 1000]

    if trajectory is None:
        print "demonstration of trajector2D_test"
        trajectory_x = np.arange(0, 3, 0.01)
        trajectory_y = np.sin(trajectory_x)
    else:
        trajectory_x = trajectory[:,0]
        trajectory_y = trajectory[:,1]

    plt.plot(trajectory_x, trajectory_y, 'r--', lw=2)

    for i, bfs in enumerate(num_bfs):
        # dmp = dmp_discrete.DMPs_discrete(dmps=2, bfs=bfs, dt=.01, goal=1, w=np.zeros((2, bfs)))
        dmp = dmp_discrete.DMPs_discrete(dmps=2, bfs=bfs)
        dmp.imitate_path(y_des=np.array([trajectory_x, trajectory_y]))
        # change the goal
        dmp.goal[0] = 1 # set the goal_x
        # dmp.goal[1] = 0.4 # set the goal_y
        dmp.y0[0] = 0.5 # set the initial x
        # dmp.y0[1] = 0.5 # set the initial y

        data_point_num = 1000 # set the point num of trajectory, now we set it to the number oftrajectory_x
        force = np.array([0,0])
        y_track, dy_track, ddy_track = dmp.rollout(tau=tau_value)

        plt.figure(2)
        plt.plot(y_track[:, 0], y_track[:, 1], lw=2)

    plt.title('DMP imitate trajectory')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.legend(["desired trajectory", "100 Bfs", "1000 Bfs"], loc='lower right')
    plt.grid()
    plt.tight_layout()
    plt.show()

def kipo_test(trajectory):
    x_org = trajectory[:, 0]
    z_org = trajectory[:, 1]
    dmp_x = dmp_discrete.DMPs_discrete(dmps=1, bfs=100)
    dmp_z = dmp_discrete.DMPs_discrete(dmps=1, bfs=100)
    dmp_x.imitate_path(y_des=x_org)
    dmp_z.imitate_path(y_des=z_org)
    data_point_num = 1000.0 # set the point num of trajectory, now we set it to the number oftrajectory_x
    tau_value = (1./ dmp_x.dt)/2000.0
    data_number = 10000
    x   = np.zeros(data_number)
    dx  = np.zeros(data_number)
    ddx = np.zeros(data_number)
    z   = np.zeros(data_number)
    dz  = np.zeros(data_number)
    ddz = np.zeros(data_number)

    # dmp_x.y0[0] = 0.1 # set the initial x
    dmp_x.goal[0] = 8 # set the goal_x
    dmp_x.reset_state()
    dmp_z.reset_state()

    for i in range(0, data_number):
        if i < 500:
            external_force = np.array([0.0])
        else:
            tau_value = (1. / dmp_x.dt) / 2000.0
            dmp_x.goal[0] = -70  # set the goal_x
            # dmp_x.goal[0] = -5  # set the goal_x
        external_force = np.array([-0.0])
        x[i], dx[i], ddx[i] = dmp_x.step(tau=tau_value, external_force=external_force)
        z[i], dz[i], ddz[i] = dmp_z.step(tau=tau_value, external_force=external_force)
        time = np.arange(data_number)
    plt.plot(time,x)
    plt.title('DMP imitate trajectory')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.legend(["desired trajectory"])
    plt.grid()
    plt.show()

# ==============================
# My Test code
# ==============================
if __name__ == "__main__":
    import load_trajectory
    # two_trajectory_test()
    trajectory = load_trajectory.loadTrajectory("swing_trajectory.txt")
    kipo_test(trajectory)
    # single_trajectory_test(trajectory[:, 1])
    # trajectory2D_test(trajectory)
